package com.su02.multi.examples.gd;

import com.su02.multi.chainrule.Function;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.tuple.Pair;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

@Slf4j
public final class GradientDescent {
    private GradientDescent() {}

    private static final ExecutorService ES = Executors.newFixedThreadPool(1000);

    @SneakyThrows
    public static Pair<Map<Integer, Double>, Double> optimize(Function f, Map<Integer, Double> start, double rate,
        double eps) {
        Map<Integer, Double> x = new HashMap<>(start);
        double lastY = f.forward(x);
        double curY;
        int iterations = 0;
        while (true) {
            List<Callable<Void>> tasks = new ArrayList<>();
            for (int lastId : x.keySet()) {
                tasks.add(() -> {
                    x.put(lastId, x.get(lastId) - rate * f.backward(lastId).forward(x));
                    return null;
                });
            }
            List<Future<Void>> futures = ES.invokeAll(tasks);

            for (Future<Void> fur : futures) {
                fur.get();
            }

            curY = f.forward(x);
            iterations++;
            if (Math.abs(curY - lastY) <= eps) {
                break;
            }
            log.info("Current Loss: {}, Delta loss: {}", Double.valueOf(curY), Double.valueOf(Math.abs(lastY - curY)));
            lastY = curY;
        }
        System.out.println("GradientDescent Iterations: " + iterations);
        return Pair.ofNonNull(x, curY);
    }
}
